"""
Mesh generation module for DLR-FEM implementation.

This module provides functionality to generate 2D rectangular meshes 
with Gauss-Lobatto points for mass-lumped finite element methods.
"""

import numpy as np
from typing import Tuple, Optional


def gauss_lobatto_points(n: int) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate Gauss-Lobatto points and weights for polynomial order n-1.
    
    Args:
        n: Number of points (polynomial order = n-1)
        
    Returns:
        Tuple of (points, weights) on [-1, 1]
    """
    if n == 2:
        # Linear case
        points = np.array([-1.0, 1.0])
        weights = np.array([1.0, 1.0])
    elif n == 3:
        # Quadratic case
        points = np.array([-1.0, 0.0, 1.0])
        weights = np.array([1/3, 4/3, 1/3])
    elif n == 4:
        # Cubic case
        points = np.array([-1.0, -np.sqrt(1/5), np.sqrt(1/5), 1.0])
        weights = np.array([1/6, 5/6, 5/6, 1/6])
    else:
        # General case using iterative method
        points = _compute_gauss_lobatto_points(n)
        weights = _compute_gauss_lobatto_weights(points)
    
    return points, weights


def _compute_gauss_lobatto_points(n: int) -> np.ndarray:
    """Compute Gauss-Lobatto points using iterative method."""
    if n < 2:
        raise ValueError("Number of points must be at least 2")
    
    # Initial guess for interior points
    x = np.cos(np.pi * np.arange(1, n-1) / (n-1))
    
    # Newton's method to find roots of derivative of Legendre polynomial
    for _ in range(10):  # Usually converges quickly
        P, Pp = _legendre_and_derivative(x, n-1)
        x = x - P / Pp
    
    # Add endpoints
    points = np.concatenate([[-1.0], x, [1.0]])
    return np.sort(points)


def _compute_gauss_lobatto_weights(points: np.ndarray) -> np.ndarray:
    """Compute Gauss-Lobatto weights."""
    n = len(points)
    weights = np.zeros(n)
    
    for i in range(n):
        P_n_minus_1 = _legendre_polynomial(points[i], n-1)
        weights[i] = 2 / ((n-1) * n * P_n_minus_1**2)
    
    return weights


def _legendre_polynomial(x: float, n: int) -> float:
    """Evaluate Legendre polynomial of degree n at x."""
    if n == 0:
        return 1.0
    elif n == 1:
        return x
    else:
        P0, P1 = 1.0, x
        for k in range(2, n+1):
            P2 = ((2*k-1) * x * P1 - (k-1) * P0) / k
            P0, P1 = P1, P2
        return P1


def _legendre_and_derivative(x: np.ndarray, n: int) -> Tuple[np.ndarray, np.ndarray]:
    """Evaluate Legendre polynomial and its derivative."""
    if n == 0:
        return np.ones_like(x), np.zeros_like(x)
    elif n == 1:
        return x, np.ones_like(x)
    
    P0 = np.ones_like(x)
    P1 = x.copy()
    Pp0 = np.zeros_like(x)
    Pp1 = np.ones_like(x)
    
    for k in range(2, n+1):
        P2 = ((2*k-1) * x * P1 - (k-1) * P0) / k
        Pp2 = ((2*k-1) * (P1 + x * Pp1) - (k-1) * Pp0) / k
        P0, P1 = P1, P2
        Pp0, Pp1 = Pp1, Pp2
    
    return P1, Pp1


class RectangularMesh:
    """
    2D rectangular mesh with Gauss-Lobatto points.
    
    This class generates a tensor product mesh using Gauss-Lobatto points
    in both x and y directions, suitable for spectral element methods.
    """
    
    def __init__(self, 
                 x_bounds: Tuple[float, float] = (0.0, 1.0),
                 y_bounds: Tuple[float, float] = (0.0, 1.0),
                 nx: int = 32,
                 ny: int = 32,
                 poly_order: int = 1):
        """
        Initialize rectangular mesh.
        
        Args:
            x_bounds: (x_min, x_max) domain bounds in x-direction
            y_bounds: (y_min, y_max) domain bounds in y-direction  
            nx: Number of grid points in x-direction
            ny: Number of grid points in y-direction
            poly_order: Polynomial order for basis functions
        """
        self.x_bounds = x_bounds
        self.y_bounds = y_bounds
        self.nx = nx
        self.ny = ny
        self.poly_order = poly_order
        
        # Generate 1D grids
        self.x_grid = np.linspace(x_bounds[0], x_bounds[1], nx)
        self.y_grid = np.linspace(y_bounds[0], y_bounds[1], ny)
        
        # Grid spacing
        self.hx = (x_bounds[1] - x_bounds[0]) / (nx - 1)
        self.hy = (y_bounds[1] - y_bounds[0]) / (ny - 1)
        
        # Create 2D coordinate meshes
        self.X, self.Y = np.meshgrid(self.x_grid, self.y_grid, indexing='ij')
        
        # Total number of degrees of freedom
        self.ndof = nx * ny
        
        # Generate Gauss-Lobatto points and weights for reference element
        self.gl_points, self.gl_weights = gauss_lobatto_points(poly_order + 1)
    
    def get_coordinates(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get 2D coordinate arrays."""
        return self.X, self.Y
    
    def get_flat_coordinates(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get flattened coordinate arrays."""
        return self.X.flatten(), self.Y.flatten()
    
    def get_grid_spacing(self) -> Tuple[float, float]:
        """Get grid spacing in x and y directions."""
        return self.hx, self.hy
    
    def map_to_reference(self, x: float, y: float, 
                        elem_x_idx: int, elem_y_idx: int) -> Tuple[float, float]:
        """
        Map physical coordinates to reference element [-1, 1]^2.
        
        Args:
            x, y: Physical coordinates
            elem_x_idx, elem_y_idx: Element indices
            
        Returns:
            (xi, eta): Reference coordinates
        """
        # Element bounds
        x_min = self.x_grid[elem_x_idx]
        x_max = self.x_grid[elem_x_idx + 1] if elem_x_idx < self.nx - 1 else x_min + self.hx
        y_min = self.y_grid[elem_y_idx]  
        y_max = self.y_grid[elem_y_idx + 1] if elem_y_idx < self.ny - 1 else y_min + self.hy
        
        # Map to [-1, 1]
        xi = 2 * (x - x_min) / (x_max - x_min) - 1
        eta = 2 * (y - y_min) / (y_max - y_min) - 1
        
        return xi, eta
    
    def jacobian(self, elem_x_idx: int, elem_y_idx: int) -> float:
        """
        Compute Jacobian determinant for element.
        
        Args:
            elem_x_idx, elem_y_idx: Element indices
            
        Returns:
            Jacobian determinant
        """
        return self.hx * self.hy / 4  # Factor of 4 from reference element mapping
    
    def get_boundary_nodes(self) -> dict:
        """
        Get indices of boundary nodes.
        
        Returns:
            Dictionary with 'left', 'right', 'bottom', 'top' boundary node indices
        """
        boundary_nodes = {
            'left': list(range(0, self.ndof, self.nx)),  # x = x_min
            'right': list(range(self.nx-1, self.ndof, self.nx)),  # x = x_max  
            'bottom': list(range(self.nx)),  # y = y_min
            'top': list(range(self.ndof - self.nx, self.ndof))  # y = y_max
        }
        return boundary_nodes
    
    def visualize_mesh(self) -> None:
        """Visualize the mesh structure."""
        import matplotlib.pyplot as plt
        
        fig, ax = plt.subplots(figsize=(8, 6))
        
        # Plot grid lines
        for i in range(self.nx):
            ax.axvline(self.x_grid[i], color='gray', alpha=0.5, linewidth=0.5)
        for j in range(self.ny):
            ax.axhline(self.y_grid[j], color='gray', alpha=0.5, linewidth=0.5)
        
        # Plot grid points
        ax.scatter(self.X.flatten(), self.Y.flatten(), c='red', s=10, alpha=0.7)
        
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        ax.set_title(f'Rectangular Mesh ({self.nx} × {self.ny} points)')
        ax.grid(True, alpha=0.3)
        ax.set_aspect('equal')
        
        plt.tight_layout()
        plt.show()


def create_uniform_mesh(domain_x: Tuple[float, float],
                       domain_y: Tuple[float, float], 
                       nx: int, ny: int) -> RectangularMesh:
    """
    Convenience function to create a uniform rectangular mesh.
    
    Args:
        domain_x: (x_min, x_max) 
        domain_y: (y_min, y_max)
        nx: Number of points in x-direction
        ny: Number of points in y-direction
        
    Returns:
        RectangularMesh instance
    """
    return RectangularMesh(domain_x, domain_y, nx, ny)


if __name__ == "__main__":
    # Test the mesh generation
    mesh = create_uniform_mesh((0, 1), (0, 1), 10, 10)
    print(f"Created mesh with {mesh.ndof} degrees of freedom")
    print(f"Grid spacing: hx = {mesh.hx:.4f}, hy = {mesh.hy:.4f}")
    
    # Test Gauss-Lobatto points
    for n in [2, 3, 4]:
        points, weights = gauss_lobatto_points(n)
        print(f"\nGauss-Lobatto points (n={n}):")
        print(f"Points: {points}")
        print(f"Weights: {weights}")
        print(f"Sum of weights: {np.sum(weights):.6f}")